package argo.hot100.heap;

import java.util.Comparator;
import java.util.HashMap;
import java.util.Map;
import java.util.PriorityQueue;

public class TopKFrequent {
    public int[] topKFrequent(int[] nums, int k) {

        int[] res = new int[k];

        if (nums == null || nums.length <= k) {
            return nums;
        }
        HashMap<Integer, Integer> map = new HashMap<>();
        for (int num : nums) {
            map.put(num, map.getOrDefault(num, 0) + 1);
        }

        PriorityQueue<int[]> queue = new PriorityQueue<>(new Comparator<int[]>() {
            @Override
            public int compare(int[] o1, int[] o2) {
                return o1[1] - o2[1];
            }
        });


        for (Map.Entry<Integer, Integer> mp : map.entrySet()) {
            Integer word = mp.getKey();
            Integer cnt = mp.getValue();
            if (queue.size() < k) {
                queue.add(new int[]{word, cnt});
            } else {
                if (queue.peek()[1] < cnt) {
                    queue.poll();
                    queue.add(new int[]{word, cnt});
                }
            }


        }

        int size = queue.size();
        for (int i = 0; i < size; i++) {
            res[i] = queue.poll()[0];
        }


        return res;
    }


}
